#include <iostream>
#include <stdlib.h>

#include "gtest/gtest.h"

#include "kpextVkContext.h"
#include "KpExtGP.h"
#include "kpextOpImgResampleU16C1toF32C1.h"

#include "kpextCtxGlobal.h"

using namespace std;

/**
 * @brief 计算生成重色散补偿采样索引，公式为$n^{'} = a n(a+1 - \frac{n}{N-1})^{-1}$，其中$a = 0.5 * (\frac{0.5}{delta} - 1)$
 * @param N  索引个数
 * @param delta  中间索引值差除以N的大小
 * @return  返回生成的缓冲区，缓冲区由调用方进行释放
 */
static float * genDCResampleIdx(int N, float delta){
    if(N <= 0){
        return nullptr;
    }
    float * ret = new float[static_cast<size_t>(N)];
    std::memset(ret, 0, static_cast<size_t>(N) * sizeof (float));
    for (int n = 0; n < N; ++n) {
        float a = 0.5f * (0.5f / delta  - 1);
        float t =  static_cast<float>(n) / static_cast<float>(N - 1);
        t = a + 1 - t;
        t = 1 / t;
        t = a * static_cast<float>(n) * t;
        ret[n] = t;
    }
    return ret;
}

void resampleBy1DIndex(int cols, int rows,
                       uint16_t * src, float * index,
                       float * dst){
    int ia, ib;
    float fv, fa, fb;
    for (int r = 0; r < rows; ++r) {
        for (int c = 0; c < cols; ++c) {
            fv = index[c];
            ia = int(fv);
            fv = fv - float(ia);
            fv = float(fv >= 0) * fv;
            ia = ia * int(ia > -1);
            ia = int(ia<cols) * ia + int(ia>=cols) * int(cols - 1);
            ib = ia + 1;
            ib = int(ib<cols) * ib + int(ib>=cols) * int(cols - 1);
            ia = int(r * cols + ia);
            ib = int(r * cols + ib);
            fa = src[ia];
            fb = src[ib];
            fa = (1-fv) * fa;
            fb = fv * fb;
            fv = fa + fb;
            dst[r * cols + c] = fv;
        }
    }
}

TEST(kpextFuncTest, OpImgResampleU16C1toF32C1){
    ASSERT_NE(nullptr, GvkContextPtr);

    int cols = 1024, rows = 1000;

    auto cmdbuf = GvkContextPtr->cmdbuf();

    size_t size = static_cast<size_t>(cols * rows);
    auto stagebuf1 = GvkContextPtr->stagingBuf(size * 2 + 4);
    auto stagebuf2 = GvkContextPtr->stagingBuf(size * 4 + 4);
    auto inbuf = GvkContextPtr->deviceBuf(size * 2 + 4);
    auto outbuf = GvkContextPtr->deviceBuf(size * 4 + 4);
    auto idxbuf = GvkContextPtr->deviceBuf(static_cast<size_t>(cols) * sizeof (float));
    auto stagebuf3 = GvkContextPtr->stagingBuf(static_cast<size_t>(cols) * sizeof (float));
    auto op = std::make_shared<kpext::OpImgResampleU16C1toF32C1>(
                GvkContextPtr
                );

    std::shared_ptr<float> result(new float[size]);

    ASSERT_NE(nullptr, cmdbuf);
    ASSERT_NE(nullptr, stagebuf1);
    ASSERT_NE(nullptr, stagebuf2);
    ASSERT_NE(nullptr, stagebuf3);
    ASSERT_NE(nullptr, inbuf);
    ASSERT_NE(nullptr, outbuf);
    ASSERT_NE(nullptr, idxbuf);
    ASSERT_NE(nullptr, op);

    std::shared_ptr<float> resampleidx(genDCResampleIdx(cols, 0.01f));
//    for (int i = 0; i < cols; ++i) {
//        resampleidx.get()[i] = static_cast<float>(i);
//    }
    {
        auto data = stagebuf3->data<float>();
        std::memcpy(data, resampleidx.get(), static_cast<size_t>(cols) * sizeof (float));
        stagebuf3->unmap();
    }

    {
        auto data = stagebuf1->data<uint16_t>();
        for (size_t i = 0; i < size; ++i) {
            data[i] = static_cast<uint16_t>(i) % static_cast<uint16_t>(cols);
        }
        resampleBy1DIndex(cols, rows, data, resampleidx.get(), result.get());
        stagebuf1->unmap();
    }

    kpext::CmdBaseBufferCpy bufcpy;
    kpext::CmdBaseBufferSet bufset;
    kpext::CmdBuffersMemBarrier barrierTx;
    kpext::CmdBuffersMemBarrier barrierRx;

    barrierTx.prepare(
                {inbuf, outbuf, idxbuf},
                vk::PipelineStageFlagBits::eTransfer,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eTransferWrite,
                vk::AccessFlagBits::eShaderRead
                );
    barrierRx.prepare(
                {inbuf, outbuf, idxbuf},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eTransfer,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eTransferRead
                );

    op->rebuild(cols, rows,
                inbuf, idxbuf, outbuf
                );

    cmdbuf->begin();
    bufcpy.prepare(inbuf, stagebuf1);
    cmdbuf->record(&bufcpy);
    bufcpy.prepare(idxbuf, stagebuf3);
    cmdbuf->record(&bufcpy);
    cmdbuf->record(&barrierTx);

    cmdbuf->record(op.get());

    bufcpy.prepare(stagebuf2, outbuf);
    cmdbuf->record(&barrierRx);
    cmdbuf->record(&bufcpy);

    cmdbuf->submit();
    cmdbuf->waitExec();

    {
        auto p1 = stagebuf1->data<uint16_t>();
        auto p2 = stagebuf2->data<float>();
            SPDLOG_INFO(__FUNCTION__ " output: [{}] ", stagebuf2->vector<float>());
        int maxarg = 0;
        float maxerr = 0;
        float maxerrv1 = 0;
        float maxerrv2 = 0;

        for (size_t i = 0; i < size; ++i) {
            float err = p2[i] - result.get()[i];
            err = fabs(err);
            if(err > maxerr){
                maxerr = err;
                maxerrv1 = p2[i];
                maxerrv2 = result.get()[i];
                maxarg = static_cast<int>(i);
            }
        }

        stagebuf1->unmap();
        stagebuf2->unmap();

        SPDLOG_INFO("max error={} at index:{},"
                    " where value pair=({:0.6f}, {:0.6f})",
                    maxerr, maxarg,
                    maxerrv1, maxerrv2);
        EXPECT_EQ(true, maxerr < 1e-6f);
    }

}

GTEST_API_ int main(int argc, char **argv) {
    printf("Running main() from %s\n", __FILE__);
    //    kpextSetLoggerLevel(KpExtLOG_LEVEL_DEBUG);
    kpextSetLoggerLevel(KpExtLOG_LEVEL_INFO);
    //  kpextSetLoggerPattern("[%D %H:%M:%S.%e] [%^%l%$] %v");
    kpextSetLoggerPattern("   [SPDLOG_%l] %v");

    testing::InitGoogleTest(&argc, argv);

    testing::Environment* const foo_env =
            testing::AddGlobalTestEnvironment(new KpExtCtxEnv);
    UNUSE_VAR(foo_env);

    return RUN_ALL_TESTS();
}
